import pyrealsense2 as rs
import numpy as np
import cv2
import os
from datetime import datetime
import open3d as o3d

# 定义一个RealSense相机类
class RealSenseCamera(object):
    def __init__(self,type):
        """
        初始化相机配置，包括分辨率和帧率。
        """
        self.type= type
        self.pipeline = rs.pipeline()
        self.config = rs.config()
        self.config.enable_stream(rs.stream.color, 1920, 1080, rs.format.bgr8, 30)
        if self.type=='l515':
            self.config.enable_stream(rs.stream.depth, 1024 , 768, rs.format.z16, 30)
        elif self.type=='d435':
            self.config.enable_stream(rs.stream.depth,  1280 ,720, rs.format.z16, 30)
        self.cfg = self.pipeline.start(self.config)
        self.profile = self.cfg.get_stream(rs.stream.color)
        self.intr = self.profile.as_video_stream_profile().get_intrinsics()

    def get_frame(self):
        """
        获取一帧彩色和深度图像，并对其进行对齐处理。
        """
        frames = self.pipeline.wait_for_frames()
        align = rs.align(rs.stream.color)
        aligned_frames = align.process(frames)
        # 获取彩色和深度帧
        color_frame = aligned_frames.get_color_frame()
        depth_frame = aligned_frames.get_depth_frame()
        color_image = np.asanyarray(color_frame.get_data())
        if self.type=='l515':
            depth_image = np.asanyarray(depth_frame.get_data())*2.5
        else:
            depth_image = np.asanyarray(depth_frame.get_data())
        return color_image,depth_image

    def __del__(self):
        """
        释放相机资源。
        """
        self.pipeline.stop()

import numpy as np
import cv2

def create_point_cloud(color_image, depth_image, camera_intr):
    """
    :param color_image: RGB 图像 (H, W, 3)
    :param depth_image: 深度图 (H, W), 每个像素代表 Z 值（单位与 scale 一致）
    :param camera_matrix: 相机内参矩阵 [fx, fy, cx, cy]
    :param depth_scale: 深度值缩放因子（例如毫米转米）
    :return: 点云数组 (N, 6), 格式为 x y z r g b
    """
    fx, fy, cx, cy = camera_intr.fx,camera_intr.fy,camera_intr.ppx,camera_intr.ppy
    h, w = depth_image.shape
    # 创建网格坐标
    u, v = np.meshgrid(np.arange(w), np.arange(h))
    z = depth_image.astype(np.float32)
    x = (u - cx) * z / fx
    y = (v - cy) * z / fy
    # 合并为三维点
    points = np.stack((x, y, z), axis=-1).reshape(-1, 3)

    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)
    pcd.colors = o3d.utility.Vector3dVector(color_image.reshape(-1, 3)/255.0)
    return pcd

if __name__ == '__main__':
    camera = RealSenseCamera('d435')

    id=0
    path=r'C:\code\camera\realsense\data\\'
    os.makedirs(path,exist_ok=True)
    while True:
        color_image,depth_image=camera.get_frame()
        # 显示图像
        cv2.imshow('Color', color_image)
        key = cv2.waitKey(1)
        if key == 32:
            timestamp=datetime.now().strftime("%Y%m%d%H%M%S")
            cv2.imwrite(path+f'depth_{timestamp}.png', (depth_image*1000).astype(np.uint16))
            cv2.imwrite(path+f'color_{timestamp}.jpg', color_image)
            print(f"保存成功")

            pcd=create_point_cloud(color_image, depth_image, camera.intr)
            o3d.visualization.draw_geometries([pcd])
            o3d.io.write_point_cloud(path+f'pointcloud_{timestamp}.ply', pcd)
        elif key == 27:
            print(f"程序终止")
            break